#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <sys/epoll.h>
#include <fcntl.h>
#include <errno.h>

#define MAX_EVENTS 1024
#define THREAD_POOL_SIZE 8
#define MAX_QUEUE_SIZE 1024
#define BUFFER_SIZE 4096

// 连接上下文结构
typedef struct {
    int fd;                 // 连接文件描述符
    char buffer[BUFFER_SIZE]; // 接收缓冲区
    size_t bytes_read;       // 已读取字节数
    int processing;          // 任务处理中标志
} connection_t;

// 任务结构
typedef struct {
    void (*task_func)(void*); // 任务处理函数
    void *arg;                // 任务参数
} task_t;

// 线程池结构
typedef struct {
    pthread_t *threads;      // 工作线程数组
    task_t *queue;           // 任务队列
    int queue_size;          // 队列大小
    int head;                // 队列头
    int tail;                // 队列尾
    int count;               // 当前任务数
    pthread_mutex_t lock;    // 队列锁
    pthread_cond_t not_empty; // 非空条件变量
    pthread_cond_t not_full;  // 非满条件变量
    int shutdown;            // 关闭标志
} thread_pool_t;

// 全局变量
thread_pool_t pool;
connection_t *connections[MAX_EVENTS] = {NULL};
pthread_mutex_t conn_lock = PTHREAD_MUTEX_INITIALIZER;

// 设置文件描述符非阻塞
void set_nonblocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

// 业务处理函数示例：数据解析和文件写入
void process_data(void *arg) {
    connection_t *conn = (connection_t *)arg;
    
    // 模拟耗时操作：数据解析
    printf("Thread %lu: Processing data from fd %d (size: %zu bytes)\n",
           pthread_self(), conn->fd, conn->bytes_read);
    
    // 模拟文件写入操作
    FILE *file = fopen("server_data.log", "a");
    if (file) {
        fwrite(conn->buffer, 1, conn->bytes_read, file);
        fclose(file);
    }
    
    // 模拟处理时间
    usleep(100000); // 100ms
    
    // 重置连接状态
    conn->bytes_read = 0;
    conn->processing = 0;
}

// 线程池工作函数
void *worker_thread(void *arg) {
    thread_pool_t *pool = (thread_pool_t *)arg;
    
    while (1) {
        pthread_mutex_lock(&pool->lock);
        
        // 等待队列非空
        while (pool->count == 0 && !pool->shutdown) {
            pthread_cond_wait(&pool->not_empty, &pool->lock);
        }
        
        if (pool->shutdown) {
            pthread_mutex_unlock(&pool->lock);
            pthread_exit(NULL);
        }
        
        // 从队列取出任务
        task_t task = pool->queue[pool->head];
        pool->head = (pool->head + 1) % pool->queue_size;
        pool->count--;
        
        // 通知可能有空间添加新任务
        pthread_cond_signal(&pool->not_full);
        pthread_mutex_unlock(&pool->lock);
        
        // 执行任务
        task.task_func(task.arg);
    }
    return NULL;
}

// 初始化线程池
void thread_pool_init(thread_pool_t *pool, int pool_size, int queue_size) {
    pool->threads = malloc(pool_size * sizeof(pthread_t));
    pool->queue = malloc(queue_size * sizeof(task_t));
    pool->queue_size = queue_size;
    pool->head = 0;
    pool->tail = 0;
    pool->count = 0;
    pool->shutdown = 0;
    
    pthread_mutex_init(&pool->lock, NULL);
    pthread_cond_init(&pool->not_empty, NULL);
    pthread_cond_init(&pool->not_full, NULL);

    int i;
    // 创建工作线程
    for (i = 0; i < pool_size; i++) {
        pthread_create(&pool->threads[i], NULL, worker_thread, pool);
    }
}

// 添加任务到线程池
int thread_pool_add_task(thread_pool_t *pool, void (*task_func)(void*), void *arg) {
    pthread_mutex_lock(&pool->lock);
    
    // 等待队列有空间,此时任务队里已经满了，1024个任务了
    while (pool->count == pool->queue_size && !pool->shutdown) {
        pthread_cond_wait(&pool->not_full, &pool->lock);
    }
    
    if (pool->shutdown) {
        pthread_mutex_unlock(&pool->lock);
        return -1;
    }
    
    // 添加任务到队列
    pool->queue[pool->tail].task_func = task_func;
    pool->queue[pool->tail].arg = arg;
    pool->tail = (pool->tail + 1) % pool->queue_size;
    pool->count++;
    
    // 通知工作线程有新任务
    pthread_cond_signal(&pool->not_empty);
    pthread_mutex_unlock(&pool->lock);
    return 0;
}

// 销毁线程池
void thread_pool_destroy(thread_pool_t *pool) {
    pthread_mutex_lock(&pool->lock);
    pool->shutdown = 1;
    pthread_cond_broadcast(&pool->not_empty);
    pthread_cond_broadcast(&pool->not_full);
    pthread_mutex_unlock(&pool->lock);

    int i;
    // 等待所有线程退出
    for (i = 0; i < THREAD_POOL_SIZE; i++) {
        pthread_join(pool->threads[i], NULL);
    }
    
    free(pool->threads);
    free(pool->queue);
    pthread_mutex_destroy(&pool->lock);
    pthread_cond_destroy(&pool->not_empty);
    pthread_cond_destroy(&pool->not_full);
}

// 创建监听socket
int create_listen_socket(int port) {
    int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_fd < 0) {
        perror("socket");
        exit(EXIT_FAILURE);
    }
    
    // 设置SO_REUSEADDR
    int opt = 1;
    setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);
    
    if (bind(listen_fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        perror("bind");
        close(listen_fd);
        exit(EXIT_FAILURE);
    }
    
    if (listen(listen_fd, 128) < 0) {
        perror("listen");
        close(listen_fd);
        exit(EXIT_FAILURE);
    }
    
    set_nonblocking(listen_fd);
    return listen_fd;
}

// 处理连接事件
void handle_connection(int epoll_fd, int listen_fd) {
    struct epoll_event ev;
    struct epoll_event events[MAX_EVENTS];
    
    while (1) {
        int nfds = epoll_wait(epoll_fd, events, MAX_EVENTS, -1);
        if (nfds == -1) {
            perror("epoll_wait");
            continue;
        }
        int i;
        for (i = 0; i < nfds; i++) {
            int fd = events[i].data.fd;
            
            // 新连接到达
            if (fd == listen_fd) {
                struct sockaddr_in client_addr;
                socklen_t addr_len = sizeof(client_addr);
                int conn_fd = accept(listen_fd, (struct sockaddr*)&client_addr, &addr_len);
                if (conn_fd < 0) {
                    perror("accept");
                    continue;
                }
                
                set_nonblocking(conn_fd);
                
                // 创建连接上下文
                connection_t *conn = malloc(sizeof(connection_t));
                conn->fd = conn_fd;
                conn->bytes_read = 0;
                conn->processing = 0;
                
                // 添加到连接表
                pthread_mutex_lock(&conn_lock);
                connections[conn_fd] = conn;
                pthread_mutex_unlock(&conn_lock);
                
                // 添加到epoll
                ev.events = EPOLLIN | EPOLLET;
                ev.data.fd = conn_fd;
                epoll_ctl(epoll_fd, EPOLL_CTL_ADD, conn_fd, &ev);
                
                printf("New connection: fd=%d\n", conn_fd);
            } 
            // 可读事件
            else if (events[i].events & EPOLLIN) {
                connection_t *conn = connections[fd];
                if (!conn) continue;
                
                // 读取数据
                ssize_t n = read(fd, conn->buffer + conn->bytes_read, 
                                 BUFFER_SIZE - conn->bytes_read);
                if (n < 0) {
                    if (errno != EAGAIN) {
                        perror("read");
                        close(fd);
                        free(conn);
                        connections[fd] = NULL;
                    }
                    continue;
                } else if (n == 0) {
                    // 客户端关闭连接
                    close(fd);
                    free(conn);
                    connections[fd] = NULL;
                    continue;
                }
                
                conn->bytes_read += n;
                
                // 检查是否收到完整数据包（这里简化处理，实际需要协议解析）
                if (conn->bytes_read > 0 && !conn->processing) {
                    conn->processing = 1;
                    
                    // 将任务添加到线程池
                    thread_pool_add_task(&pool, process_data, conn);
                }
            }
        }
    }
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        fprintf(stderr, "Usage: %s <port>\n", argv[0]);
        exit(EXIT_FAILURE);
    }
    
    int port = atoi(argv[1]);
    int listen_fd = create_listen_socket(port);
    
    // 初始化线程池
    thread_pool_init(&pool, THREAD_POOL_SIZE, MAX_QUEUE_SIZE);
    
    // 创建epoll实例
    int epoll_fd = epoll_create1(0);
    if (epoll_fd < 0) {
        perror("epoll_create1");
        exit(EXIT_FAILURE);
    }
    
    // 添加监听socket到epoll
    struct epoll_event ev;
    ev.events = EPOLLIN;
    ev.data.fd = listen_fd;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listen_fd, &ev) < 0) {
        perror("epoll_ctl: listen_sock");
        exit(EXIT_FAILURE);
    }
    
    printf("Server started on port %d\n", port);
    printf("Using thread pool with %d threads\n", THREAD_POOL_SIZE);
    
    // 主事件循环
    handle_connection(epoll_fd, listen_fd);
    
    // 清理资源（正常情况下不会执行到这里）
    close(epoll_fd);
    close(listen_fd);
    thread_pool_destroy(&pool);
    return 0;
}

